/*
 * Copyright(C) 2020. Huawei Technologies Co.,Ltd. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include "FaceLandmark.h"
#include <fstream>
#include <iostream>
#include <memory>
#include <cmath>
#include <CostStatistic/CostStatistic.h>

#include "opencv2/highgui.hpp"
#include "opencv2/imgcodecs.hpp"
#include "opencv2/imgproc.hpp"

#include "Log/Log.h"
#include "FileEx/FileEx.h"
#include "PointerDeleter/PointerDeleter.h"
#include "TestCV/TestCV.h"

namespace ascendFaceRecognition {
namespace {
const uint16_t DEGREE90 = 90;
const float ELEMENT_POSITION_LIMIT = 48.0;
const int FACE_KEY_POINT_SIZE = 5;
} // namespace
FaceLandmark::FaceLandmark()
{
    isStop_ = true;
    instanceId_ = -1;
}

FaceLandmark::~FaceLandmark() {}

APP_ERROR FaceLandmark::InitResource(void)
{
    LogDebug << "FaceLandmark[" << instanceId_ << "]: Begin to init device resource.";
    size_t inputNum = modelInfer_->GetModelNumInputs();
    if (inputNum != 1) {
        LogFatal << "input tensor size is invaild" << inputNum;
        return APP_ERR_COMM_FAILURE;
    }
    LogDebug << "inputNum = " << inputNum;
    APP_ERROR ret;
    for (size_t i = 0; i < inputNum; i++) {
        void* buffer = nullptr;
        size_t size = modelInfer_->GetModelInputSizeByIndex(i);
        inputDataSet_.bufs.push_back(buffer);
        inputDataSet_.sizes.push_back(size);
        LogDebug << "FaceLandmark[" << instanceId_ << "]: model input tensor i = " << i << ", size= " << size;
    }

    size_t outputNum = modelInfer_->GetModelNumOutputs();
    if (outputNum == 0) {
        LogFatal << "FaceLandmark[" << instanceId_ << "]: output tensor size is invaild " << outputNum;
        return APP_ERR_COMM_FAILURE;
    }

    LogDebug << "FaceLandmark[" << instanceId_ << "]: outputNum = " << outputNum;
    for (size_t i = 0; i < outputNum; i++) {
        void* buffer = nullptr;
        size_t size = modelInfer_->GetModelOutputSizeByIndex(i);
        ret = aclrtMalloc(&buffer, size, ACL_MEM_MALLOC_HUGE_FIRST);
        if (ret != APP_ERR_OK) {
            LogFatal << "FaceLandmark[" << instanceId_ << "]: create output failed!";
            return ret;
        }
        outputDataSet_.bufs.push_back(buffer);
        outputDataSet_.sizes.push_back(size);
        LogDebug << "FaceLandmark[" << instanceId_ << "]: model output tensor i = " << i << ", size = " << size;

#ifndef ASCEND_ACL_OPEN_VESION
        auto bufferHost = std::shared_ptr<uint8_t>();
        bufferHost.reset(new uint8_t[size], std::default_delete<uint8_t[]>());
        if (bufferHost == nullptr) {
            LogFatal << "FaceLandmark[" << instanceId_ << "]: create output failed!";
            return APP_ERR_COMM_ALLOC_MEM;
        }
        outputBuffersHost_.push_back(bufferHost);
#endif
    }

    return APP_ERR_OK;
}

APP_ERROR FaceLandmark::ParseConfig(ConfigParser &configParser)
{
    LogDebug << "FaceLandmark[" << instanceId_ << "]: begin to parse config values.";
    std::string itemCfgStr;
    std::vector<APP_ERROR> vecRet;
    const uint8_t itemCfgMaxNum = 64;
    vecRet.resize(itemCfgMaxNum);
    uint8_t indexVecRet = 0;

    itemCfgStr = moduleName_ + "." + pipelineName_ + std::string(".batch_size");
    vecRet[indexVecRet++] = configParser.GetUnsignedIntValue(itemCfgStr, batchSize_);

    itemCfgStr = moduleName_ + "." + pipelineName_ + std::string(".model_path");
    vecRet[indexVecRet++] = configParser.GetStringValue(itemCfgStr, modelPath_);

    itemCfgStr = std::string("max_face_num_per_frame");
    vecRet[indexVecRet++] = configParser.GetUnsignedIntValue(itemCfgStr, maxFaceNumPerFrame_);

    itemCfgStr = moduleName_ + std::string(".heatmapWidth");
    vecRet[indexVecRet++] = configParser.GetIntValue(itemCfgStr, heatMapWidth_);

    itemCfgStr = moduleName_ + std::string(".heatmapHeight");
    vecRet[indexVecRet++] = configParser.GetIntValue(itemCfgStr, heatMapHeight_);

    itemCfgStr = moduleName_ + std::string(".keyPointWeight");
    vecRet[indexVecRet++] = configParser.GetFloatValue(itemCfgStr, keyPointWeight_);

    itemCfgStr = moduleName_ + std::string(".eulerWeight");
    vecRet[indexVecRet++] = configParser.GetFloatValue(itemCfgStr, eulerWeight_);

    itemCfgStr = moduleName_ + std::string(".bigFaceWeight");
    vecRet[indexVecRet++] = configParser.GetFloatValue(itemCfgStr, bigFaceWeight_);

    itemCfgStr = moduleName_ + std::string(".minOutputScore");
    vecRet[indexVecRet++] = configParser.GetFloatValue(itemCfgStr, minOutputScore_);

    itemCfgStr = moduleName_ + std::string(".registModeEnable");
    vecRet[indexVecRet++] = configParser.GetIntValue(itemCfgStr, registModeEnable_);

    for (uint8_t i = 0; i < indexVecRet; ++i) {
        if (vecRet[i] != APP_ERR_OK) {
            return vecRet[i];
        }
    }

    LogDebug << "FaceLandmark[" << instanceId_ << "]" << " batchSize_:" << batchSize_ << " modelPath_:"
             << modelPath_.c_str() << " maxFaceNumPerFrame_:" << maxFaceNumPerFrame_ << " heatmapWidth_:"
             << heatMapWidth_ << " heatmapHeight:" << heatMapHeight_ << " registModeEnable_:" << registModeEnable_;

    return APP_ERR_OK;
}

APP_ERROR FaceLandmark::Init(ConfigParser &configParser, ModuleInitArgs &initArgs)
{
    LogDebug << "FaceLandmark[" << instanceId_ << "]: Begin to init face landmark instance" << initArgs.instanceId;

    APP_ERROR ret = aclrtCreateStream(&mdlStream_);
    if (ret != APP_ERR_OK) {
        LogError << "FaceLandmark[" << instanceId_ << "]: aclrtCreateStream failed, ret=" << ret << ".";
        return ret;
    }

    AssignInitArgs(initArgs);

    isStop_ = false;

    // initialize config params
    ret = ParseConfig(configParser);
    if (ret != APP_ERR_OK) {
        LogFatal << "FaceLandmark[" << instanceId_ << "]: Fail to parse config params, ret=" << ret << "(" <<
                 GetAppErrCodeInfo(ret) << ").";
        return ret;
    }

    if (maxFaceNumPerFrame_ < 1) {
        LogFatal << "FaceLandmark[" << instanceId_ << "]: max_face_num_per_frame must be larger than 0";
        return APP_ERR_COMM_FAILURE;
    }

    // init and get model
    modelInfer_ = ModelResource::GetInstance().GetModelProcess(modelPath_, instanceId_);
    if (modelInfer_ == nullptr) {
        LogFatal << "FaceLandmark[" << instanceId_ << "]: init model failed";
        return ret;
    }
    // init resource
    ret = InitResource();
    if (ret != APP_ERR_OK) {
        LogFatal << "FaceLandmark[" << instanceId_ << "]: init resource failed";
        return ret;
    }

    return APP_ERR_OK;
}

APP_ERROR FaceLandmark::DeInit(void)
{
    LogDebug << "FaceLandmark[" << instanceId_ << "]::begin to deinit.";

    if (mdlStream_) {
        APP_ERROR ret = aclrtDestroyStream(mdlStream_);
        if (ret != APP_ERR_OK) {
            LogError << "Failed to destroy stream, ret = " << ret;
        }
        mdlStream_ = nullptr;
    }

    for (uint32_t i = 0; i < outputDataSet_.bufs.size(); i++) {
        acldvppFree(outputDataSet_.bufs[i]);
        outputDataSet_.bufs[i] = nullptr;
    }
    for (uint32_t i = 0; i < inputDataSet_.bufs.size(); i++) {
        inputDataSet_.bufs[i] = nullptr;
    }
    return APP_ERR_OK;
}

APP_ERROR FaceLandmark::ParseResult(FaceObject &faceObject, uint32_t pos)
{
    if (batchSize_ == 0) {
        LogError << "FaceLandmark[" << instanceId_ << "]: batchSize_ is zero.";
        return APP_ERR_ACL_INVALID_PARAM;
    }
    size_t eulerBatchLen = outputDataSet_.sizes[0] / batchSize_;
    size_t keyPointBatchLen = outputDataSet_.sizes[1] / batchSize_;

    uint16_t indexEulerPtr = 0;
#ifdef ASCEND_ACL_OPEN_VESION
    float *eulerPtr = (float *)((uint8_t *)outputDataSet_.bufs[0] + pos * eulerBatchLen);
    float *heatmapPtr = (float *)((uint8_t *)outputDataSet_.bufs[1] + pos * keyPointBatchLen);
#else
    auto* eulerPtr = (float*) ((uint8_t*) outputBuffersHost_[0].get() + pos * eulerBatchLen);
    auto* heatmapPtr = (float*) ((uint8_t*) outputBuffersHost_[1].get() + pos * keyPointBatchLen);
#endif
    faceObject.faceQuality.angleYaw = fabs(eulerPtr[indexEulerPtr++]) * DEGREE90;
    faceObject.faceQuality.anglePitch = fabs(eulerPtr[indexEulerPtr++]) * DEGREE90;
    faceObject.faceQuality.angleRoll = fabs(eulerPtr[indexEulerPtr++]) * DEGREE90;

    for (int i = 0; i < heatMapHeight_; i++) {
        float* tmpPtr = heatmapPtr + i * heatMapWidth_;
        float tmpScore = *std::max_element(tmpPtr, tmpPtr + heatMapWidth_);
        faceObject.faceQuality.keyPoints.push_back(tmpScore);
    }

    const uint8_t doubleH = 2;
    uint32_t landmarkLen = heatMapHeight_ * doubleH * sizeof(float);
    auto landmarks = std::make_shared<uint8_t>();
    landmarks.reset(new uint8_t[landmarkLen], std::default_delete<uint8_t[]>());
    uint8_t* landmarkPtr = landmarks.get();
    for (int i = 0; i < heatMapHeight_; i++) {
        float* tmpPtr = heatmapPtr + i * heatMapWidth_;
        int position = std::max_element(tmpPtr, tmpPtr + heatMapWidth_) - tmpPtr;
        auto x = static_cast<float>((position % uint32_t(ELEMENT_POSITION_LIMIT)) / ELEMENT_POSITION_LIMIT);
        std::copy((uint8_t*) &x, (uint8_t*) &x + sizeof(float), landmarkPtr);

        landmarkPtr += sizeof(float);
        auto y = static_cast<float>((position / ELEMENT_POSITION_LIMIT) / ELEMENT_POSITION_LIMIT);
        std::copy((uint8_t*) &y, (uint8_t*) &y + sizeof(float), landmarkPtr);
        landmarkPtr += sizeof(float);
    }
#ifdef ASCEND_ACL_OPEN_VESION
    faceObject.landmarks.deviceData = landmarks;
#else
    faceObject.landmarks.hostData = landmarks;
#endif
    faceObject.landmarks.dataSize = landmarkLen;

    return APP_ERR_OK;
}

void FaceLandmark::SendData(std::shared_ptr<FrameMemoryInfo> &frameMemoryInfo)
{
    for (size_t i = 0; i < frameMemoryInfo->face.size(); i++) {
        float score = frameMemoryInfo->face[i].faceQuality.score;
        if (score < minOutputScore_) {
            LogError << "FaceLandmark[" << instanceId_ << "]: score is too low, score=" << score;
            continue;
        }
        std::shared_ptr<FaceObject> faceObject = std::make_shared<FaceObject>(frameMemoryInfo->face[i]);
        if (registModeEnable_ == 0) {
            SendToNextModule(MT_WARP_AFFINE, faceObject, faceObject->frameInfo.channelId);
            continue;
        }

        if (registFaceObject_.get() == nullptr || ((registFaceObject_.get() != nullptr) &&
           (registFaceObject_->faceQuality.score < frameMemoryInfo->face[i].faceQuality.score))) {
            registFaceObject_ = faceObject;
        }
    }

    if (registModeEnable_ == 1) {
        if (registFaceObject_ != nullptr) {
            SendToNextModule(MT_WARP_AFFINE, registFaceObject_, registFaceObject_->frameInfo.channelId);
            registFaceObject_.reset();
        }
    }
}

APP_ERROR FaceLandmark::PostData(std::shared_ptr<FrameMemoryInfo> &frameMemoryInfo)
{
#ifndef ASCEND_ACL_OPEN_VESION
    for (size_t i = 0; i < outputBuffersHost_.size(); i++) {
        APP_ERROR ret = aclrtMemcpy(outputBuffersHost_[i].get(), outputDataSet_.sizes[i], outputDataSet_.bufs[i],
                                    outputDataSet_.sizes[i], ACL_MEMCPY_DEVICE_TO_HOST);
        if (ret != APP_ERR_OK) {
            LogError << "FaceLandmark[" << instanceId_ << "]: Fail to copy memory.";
            return ret;
        }
    }
#endif
    for (size_t i = 0; i < frameMemoryInfo->face.size(); i++) {
        APP_ERROR ret = ParseResult(frameMemoryInfo->face[i], i);
        if (ret != APP_ERR_OK) {
            return ret;
        }
    }

    return APP_ERR_OK;
}

APP_ERROR FaceLandmark::ProcessEmptyFaceFrame(std::shared_ptr<FrameMemoryInfo> &frameMemoryInfo)
{
    if (frameMemoryInfo->infoVector[0].mode == FRAME_MODE_SEARCH) {
        SendToNextModule(MT_MOT_EMBEDDING, frameMemoryInfo, frameMemoryInfo->infoVector[0].channelId);
    } else {
        SendToNextModule(MT_WARP_AFFINE, nullptr, frameMemoryInfo->infoVector[0].channelId);
    }
    return APP_ERR_OK;
}

APP_ERROR FaceLandmark::ProcessMulitFaceFrame(std::shared_ptr<FrameMemoryInfo> &frameMemoryInfo)
{
    APP_ERROR ret;
    // model inference
    if (inputDataSet_.sizes[0] != frameMemoryInfo->landmarkMemoryBlockSize) {
        LogError << "landmarkMemoryBlockSize is not equal to the inputSizes," << "inputSizes = "
                 << inputDataSet_.sizes[0] << " landmarkMemoryBlockSize = " << frameMemoryInfo->landmarkMemoryBlockSize;
        return APP_ERR_COMM_INVALID_PARAM;
    }

    inputDataSet_.bufs[0] = frameMemoryInfo->landmarkMemoryBlock.get();
    ret = modelInfer_->ModelInference(inputDataSet_, outputDataSet_, mdlStream_);
    if (ret != APP_ERR_OK) {
        frameMemoryInfo->landmarkMemoryBlock.reset();
        inputDataSet_.bufs[0] = nullptr;
        return ret;
    }
    inputDataSet_.bufs[0] = nullptr;
    frameMemoryInfo->landmarkMemoryBlock.reset();
    ret = PostData(frameMemoryInfo);
    if (ret != APP_ERR_OK) {
        return ret;
    }

    if (frameMemoryInfo->infoVector[0].mode == FRAME_MODE_SEARCH) {
        LogDebug << "FaceLandmark Send " << frameMemoryInfo->infoVector[0].channelId;
        SendToNextModule(MT_MOT_EMBEDDING, frameMemoryInfo, frameMemoryInfo->infoVector[0].channelId);
    } else {
        CalTotalScore(frameMemoryInfo);
        SendData(frameMemoryInfo);
    }

    return APP_ERR_OK;
}

APP_ERROR FaceLandmark::Process(std::shared_ptr<void> inputData)
{
    LogDebug << "FaceLandmark[" << instanceId_ << "]: Begin to process data.";
    std::shared_ptr<FrameMemoryInfo> frameMemoryInfo = std::static_pointer_cast<FrameMemoryInfo>(inputData);
    LogDebug << "FaceLandmark[" << instanceId_ << "]: Recv " << frameMemoryInfo->infoVector[0].channelId;
    if (frameMemoryInfo == nullptr) {
        LogError << "FrameCache FaceLandmark[" << instanceId_ << "]: can't Found FrameAiInfo!";
        return APP_ERR_COMM_FAILURE;
    }
    APP_ERROR ret;
    if (frameMemoryInfo->face.empty()) {
        ret = ProcessEmptyFaceFrame(frameMemoryInfo);
    } else {
        ret = ProcessMulitFaceFrame(frameMemoryInfo);
    }
    return ret;
}

void FaceLandmark::CalTotalScore(std::shared_ptr<FrameMemoryInfo> &frameMemoryInfo)
{
    for (size_t i = 0; i < frameMemoryInfo->face.size(); ++i) {
        float keyPointScore = CalKeyPointScore(frameMemoryInfo->face[i]); // calculate point score
        float eulerScore = CalEulerScore(frameMemoryInfo->face[i]);
        float bigFaceScore = CalBigFaceScore(frameMemoryInfo->face[i]);
        float score = keyPointWeight_ * keyPointScore + eulerWeight_ * eulerScore + bigFaceWeight_ * bigFaceScore;
        frameMemoryInfo->face[i].faceQuality.score = score;
    }
}

float FaceLandmark::CalKeyPointScore(const FaceObject &faceObject)
{
    float score = 0;
    if (faceObject.faceQuality.keyPoints.empty()) {
        LogWarn << "FaceLandmark[" << instanceId_ << "]: keyPoints size is zero.";
        return score;
    }

    if (faceObject.faceQuality.keyPoints.size() != FACE_KEY_POINT_SIZE) {
        LogWarn << "FaceLandmark[" << instanceId_ << "]: keyPoints size is invalid.";
        return score;
    }
    const float elementScoreLimit = 0.2;
    for (int i = 0; i < FACE_KEY_POINT_SIZE; ++i) {
        float tmpScore = faceObject.faceQuality.keyPoints[i];
        score += ((tmpScore > elementScoreLimit) ? elementScoreLimit : tmpScore);
    }
    return score;
}

float FaceLandmark::CalEulerScore(const FaceObject &faceObject)
{
    const uint32_t degree90 = 90;
    float yaw = faceObject.faceQuality.angleYaw;
    float pitch = faceObject.faceQuality.anglePitch;
    float roll = faceObject.faceQuality.angleRoll;
    const uint32_t pitchConstant = 6;
    pitch = (pitch > pitchConstant) ? (pitch - pitchConstant) : 0;
    return (degree90 - yaw) / degree90 + (degree90 - pitch) / degree90 + (degree90 - roll) / degree90;
}

float FaceLandmark::CalBigFaceScore(const FaceObject &faceObject)
{
    float width = faceObject.info.width;
    float height = faceObject.info.height;
    const uint16_t maxFaceHW = 60;
    const uint16_t normFaceHW = 50;
    const float faceStretchRatio = 1.2;
    const float faceScoreConstant = 3600.0;

    width = (width > normFaceHW) ? maxFaceHW : (width * faceStretchRatio);
    height = (height > normFaceHW) ? maxFaceHW : (height * faceStretchRatio);
    return 1 - fabsf(maxFaceHW - width) * fabsf(maxFaceHW - height) / faceScoreConstant;
}
} // namespace ascendFaceRecognition
